import pandas as pd
import os
import numpy as np
from tqdm import tqdm
from talib.abstract import *
from feature_utils import *
import sys
import argparse
import datetime
import random
import json

date_range = [5, 13]
SEQ_DAYS = [0, 2, 5, 7]

def bottom_top(df, index, row):
    if index < 150:
        return None
    c1 = row.close
    h1 = row.high
    l1 = row.low
    # pos
    # down x up
    # -    x up
    tflag = True
    lflag = True
    target = None

    hcloses = []
    lcloses = []
    for sqd in range(1, 5):
        if df.loc[index + sqd, 'low'] > l1 and df.loc[index - sqd, 'low'] > l1:
            pass
        else:
            lflag = False
            break

    for sqd in range(1, 5):
        if df.loc[index + sqd, 'high'] < h1:
            pass
        else:
            tflag = False
            break

    lastc9a = df.loc[index + 9, 'close']
    lastc20b = c1 / df.loc[index - 20, 'close']
    lastc60b = c1 / df.loc[index - 60, 'close']
    lastc120b = c1 / df.loc[index - 120, 'close']


    if lflag == True and (c1 * 1.06) < lastc9a and lastc20b < 1.2 and lastc60b < 1.4 and lastc120b < 1.6:
        target = 1.0
    elif tflag == True and (c1 * 0.97) > lastc9a:
            #and random.random() < 0.3: # negative downsampling
        target = 0.0

    #if target is not None:
    #    print(row.date, target, lastc20b, lastc60b, lastc120b)
    return target





def calc_rsi(df):
    features = []
    for d in [6, 12, 24]:
        dfn = 'rsi_' + str(d)
        features.append(dfn)
        df[dfn] = RSI(df, timeperiod = d) / 100
    return df, features

def calc_bbands(df):
    features = ['bbp', "b_up_low"]
    bands = BBANDS(df, timeperiod = 20, nbdevup = 2.0, nbdevdn = 2.0, matype = 0)
    df['bbp'] = (df.close - bands.lowerband) / (bands.upperband - bands.lowerband)
    df['b_up_low'] = bands.upperband / bands.lowerband
    return df, features

def calc_macd(df):
    macd_df = MACD(df, fastperiod=12, slowperiod=26, signalperiod=9)
    # macd  macdsignal  macdhist
    df = pd.concat([df, macd_df], axis = 1)
    macd_features = ['macd','macdsignal','macdhist']
    return df, macd_features

def calc_kdj(df):
    matype = 0
    # slowk      slowd
    stoch = STOCH(df, fastk_period=9, slowk_matype=matype, slowk_period=3, slowd_period=3, slowd_matype=matype)
    df['stoch_k'] = stoch['slowk'] / 100
    df['stoch_d'] = stoch['slowd'] / 100
    df['stoch_j'] = (stoch['slowk'] * 3 - stoch['slowd'] * 2)/100
    kdj_features = ['stoch_k', 'stoch_d', 'stoch_j']
    return df, kdj_features


def calc_close_ma_diff(df):
    sma_close_diff_names = []
    for d in date_range:
        diff_name = 'SMA_DIFF_' + str(d)
        sma_name = 'SMA_' + str(d)
        sma_close_diff_names.append(diff_name)

        df[sma_name] = SMA(df, timeperiod = d)
        df[diff_name] = (df['close'] - df[sma_name]) / df[sma_name]

    return df, sma_close_diff_names


def calc_vol_ma_diff(df):
    sma_vol_diff_names = []
    for d in date_range:
        diff_name = 'VOL_SMA_DIFF_' + str(d)
        sma_name = 'VOL_SMA_' + str(d)
        sma_vol_diff_names.append(diff_name)

        df[sma_name] = SMA(df, timeperiod = d, price = 'volume')
        df[diff_name] = (df['volume'] - df[sma_name]) / df[sma_name]
    return df, sma_vol_diff_names


def calc_stats(df):
    features = []
    features.append("BETA")
    df["BETA"] = BETA(df.high, df.low, timeperiod = 5)

    features.append("CORREL")
    df["CORREL"] = CORREL(df.high, df.low, timeperiod=30)

    # features.append("LINEARREG")
    # df["LINEARREG"] = LINEARREG(df.close, timeperiod=14)

    # features.append("LINEARREG_ANGLE")
    # df["LINEARREG_ANGLE"] = LINEARREG_ANGLE(df.close, timeperiod=14)

    # features.append("LINEARREG_INTERCEPT")
    # df["LINEARREG_INTERCEPT"] = LINEARREG_INTERCEPT(df.close, timeperiod=14)

    # features.append("LINEARREG_SLOPE")
    # df["LINEARREG_SLOPE"] = LINEARREG_SLOPE(df.close, timeperiod=14)

    features.append("STDDEV")
    df["STDDEV"] = STDDEV(df.close, timeperiod=5)

    # print(df[features].tail(10))

    return df, features



def calc_cdl(df):
    features = []
    for cdln in CDL_NAMES:
        cdln_func = eval(cdln)
        df[cdln] = (cdln_func(df) > 0).astype(int)
        features.append(cdln)
    return df, features

def calc_cci(df):
    df["CCI"] = CCI(df.high, df.low, df.close, timeperiod=14) / 100
    df["WILLR"] = WILLR(df.high, df.low, df.close, timeperiod=14) /100
    # print(df[["CCI", "WILLR"]].tail(10))
    return df, ["CCI", "WILLR"]

def extract_feature(df, index):
    # ma5 diff
    return df.loc[index, INF_COLUMNS]


def add_extra_indicators(df):
    features = []

    indicator_funcs = [
        calc_rsi,
        calc_kdj,
        calc_bbands,
        calc_macd,
        calc_close_ma_diff,
        calc_vol_ma_diff,
        calc_stats,
        calc_cci,
        # calc_cdl
    ]
    for indif in indicator_funcs:
        df, fns = indif(df)
        features += fns
    #print(features)
    df = df.fillna(method="bfill").fillna(method="ffill").fillna(0)
    df.replace([np.inf, -np.inf], 0, inplace = True)
    return df
